package main

import (
	"flag"
	"log"
	"net/http"
	"os/exec"

	"github.com/creack/pty"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// 升级 WebSocket 连接的配置
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 在生产环境中应严格检查 Origin
	},
}

func handleWebSocket(c *gin.Context) {
	// 升级 HTTP 连接到 WebSocket
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return
	}
	defer conn.Close()

	// 启动一个 shell (Unix-like 系统使用 "sh", Windows 可改为 "cmd")
	cmd := exec.Command("sh", "-c", "TERM=xterm-256color /bin/bash") // 根据系统调整

	// 创建伪终端
	ptty, err := pty.Start(cmd)
	if err != nil {
		log.Printf("PTY 启动失败: %v", err)
		conn.WriteMessage(websocket.TextMessage, []byte("\r\nError: 无法启动终端会话\r\n"))
		return
	}
	defer func() {
		cmd.Process.Kill()
		cmd.Process.Wait()
		ptty.Close()
	}()

	// 将 PTY 的输出转发到 WebSocket
	go func() {
		buffer := make([]byte, 1024)
		for {
			n, err := ptty.Read(buffer)
			if err != nil {
				log.Printf("从 PTY 读取失败: %v", err)
				return
			}
			err = conn.WriteMessage(websocket.TextMessage, buffer[:n])
			if err != nil {
				log.Printf("写入 WebSocket 失败: %v", err)
				return
			}
		}
	}()

	// 从 WebSocket 读取数据（用户输入）并写入 PTY
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			log.Printf("从 WebSocket 读取失败: %v", err)
			break
		}
		_, err = ptty.Write(message)
		if err != nil {
			log.Printf("写入 PTY 失败: %v", err)
			break
		}
	}
}

func main() {
	// 解析命令行参数
	port := flag.String("port", "8080", "指定服务器监听的端口")
	flag.Parse()

	r := gin.Default()

	// 提供静态文件服务 (HTML, CSS, JS)
	r.StaticFS("/static", http.Dir("./static"))

	// WebSocket 端点
	r.GET("/ws", handleWebSocket)

	// 主页面
	r.GET("/", func(c *gin.Context) {
		c.File("./static/index.html")
	})

	log.Printf("服务器启动在 :%s 端口\n", *port)
	log.Fatal(r.Run(":" + *port))
}
